#include <QString>
#include <qDebug>

#include <XProcess>

#include "Constant.h"
#include "ProcessHooker.h"

using namespace XProcess;

BOOL
WINAPI
hook_CreateProcessW(decltype(&CreateProcessW) real_CreateProcessW,
_In_opt_ LPCWSTR lpApplicationName,
_Inout_opt_ LPWSTR lpCommandLine,
_In_opt_ LPSECURITY_ATTRIBUTES lpProcessAttributes,
_In_opt_ LPSECURITY_ATTRIBUTES lpThreadAttributes,
_In_ BOOL bInheritHandles,
_In_ DWORD dwCreationFlags,
_In_opt_ LPVOID lpEnvironment,
_In_opt_ LPCWSTR lpCurrentDirectory,
_In_ LPSTARTUPINFOW lpStartupInfo,
_Out_ LPPROCESS_INFORMATION lpProcessInformation
)
{
    auto appname = QString::fromUtf16((ushort *)lpApplicationName);
    auto cmdline = QString::fromUtf16((ushort *)lpCommandLine);
    qDebug() << "hook_CreateProcessW : " << appname + "[" + cmdline + "]" << endl;
    auto retv = real_CreateProcessW(lpApplicationName,
        lpCommandLine,
        lpProcessAttributes,
        lpThreadAttributes,
        bInheritHandles,
        dwCreationFlags | CREATE_SUSPENDED,
        lpEnvironment,
        lpCurrentDirectory,
        lpStartupInfo,
        lpProcessInformation);
    XProcess::Process::setDebug(true);
    XProcess::Thread thread;
    XProcess::Process::fromHandleTakeOver(lpProcessInformation->hProcess, OpenMode::ReadAndWrite).injectDllAndRun("InjectedDll", "start", nullptr, thread);
    thread.wait();
    ::ResumeThread(lpProcessInformation->hThread);
    return retv;
}

BOOL
WINAPI
hook_CreateProcessA(decltype(&CreateProcessA) real_CreateProcessA,
_In_opt_ LPCSTR lpApplicationName,
_Inout_opt_ LPSTR lpCommandLine,
_In_opt_ LPSECURITY_ATTRIBUTES lpProcessAttributes,
_In_opt_ LPSECURITY_ATTRIBUTES lpThreadAttributes,
_In_ BOOL bInheritHandles,
_In_ DWORD dwCreationFlags,
_In_opt_ LPVOID lpEnvironment,
_In_opt_ LPCSTR lpCurrentDirectory,
_In_ LPSTARTUPINFOA lpStartupInfo,
_Out_ LPPROCESS_INFORMATION lpProcessInformation
)
{
    auto appname = QString::fromLocal8Bit(lpApplicationName);
    auto cmdline = QString::fromLocal8Bit(lpCommandLine);
    qDebug() << "hook_CreateProcessW : " << appname + "[" + cmdline + "]" << endl;
    auto retv = real_CreateProcessA(lpApplicationName,
        lpCommandLine,
        lpProcessAttributes,
        lpThreadAttributes,
        bInheritHandles,
        dwCreationFlags | CREATE_SUSPENDED,
        lpEnvironment,
        lpCurrentDirectory,
        lpStartupInfo,
        lpProcessInformation);
    XProcess::Process::setDebug(true);
    XProcess::Thread thread;
    XProcess::Process::fromHandleTakeOver(lpProcessInformation->hProcess, OpenMode::ReadAndWrite).injectDllAndRun("InjectedDll", "start", nullptr, thread);
    thread.wait();
    ::ResumeThread(lpProcessInformation->hThread);
    return retv;

}
XProcess::JmpHook<decltype(&CreateProcessW)> h_CreateProcessW;
XProcess::JmpHook<decltype(&CreateProcessA)> h_CreateProcessA;
void ProcessHooker::hookAll()
{
    h_CreateProcessW = XProcess::JmpHooker<(int)Constants::HookIds::CreateProcessW, decltype(&CreateProcessW)>::hook(&CreateProcessW, &hook_CreateProcessW);
    h_CreateProcessA = XProcess::JmpHooker<(int)Constants::HookIds::CreateProcessA, decltype(&CreateProcessA)>::hook(&CreateProcessA, &hook_CreateProcessA);
}

void ProcessHooker::unhookAll()
{
    h_CreateProcessW.unHook();
    h_CreateProcessA.unHook();
}
